######################################################################################
#U-Net: Convolutional Networks for BiomedicalImage Segmentation
#Paper-Link: https://arxiv.org/pdf/1505.04597.pdf
######################################################################################

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary



__all__ = ["MANet"]


class double_conv(nn.Module):
    '''(conv => BN => ReLU) * 2'''

    def __init__(self, in_ch, out_ch):
        super(double_conv, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.conv(x)
        return x


class inconv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(inconv, self).__init__()
        self.conv = double_conv(in_ch, out_ch)

    def forward(self, x):
        x = self.conv(x)
        return x


class down(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(down, self).__init__()
        self.mpconv = nn.Sequential(
            nn.MaxPool2d(2),
            double_conv(in_ch, out_ch)
        )

    def forward(self, x):
        x = self.mpconv(x)
        return x


class up1(nn.Module):
    def __init__(self, in_ch, out_ch, bilinear=False):
        super(up1, self).__init__()
        self.bilinear = bilinear

        self.up = nn.ConvTranspose2d(in_ch, out_ch, 2, stride=2)

        self.conv = double_conv(in_ch+out_ch, out_ch)

    def forward(self, x1, x2):
        if self.bilinear:
            x1 = F.interpolate(x1, scale_factor=2, mode='bilinear', align_corners=True)
        else:
            x1 = self.up(x1)

        # input is CHW
        diffY = x2.size()[2] - x1.size()[2]
        diffX = x2.size()[3] - x1.size()[3]

        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
                        diffY // 2, diffY - diffY // 2])

        # for padding issues, see
        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd

        x = torch.cat([x2, x1], dim=1)
        x = self.conv(x)
        return x


class up(nn.Module):
    def __init__(self, in_ch, out_ch, bilinear=False):
        super(up, self).__init__()
        self.bilinear = bilinear

        self.up = nn.ConvTranspose2d(in_ch, out_ch, 2, stride=2)

        self.conv = double_conv(in_ch, out_ch)

    def forward(self, x1, x2):
        if self.bilinear:
            x1 = F.interpolate(x1, scale_factor=2, mode='bilinear', align_corners=True)
        else:
            x1 = self.up(x1)

        # input is CHW
        diffY = x2.size()[2] - x1.size()[2]
        diffX = x2.size()[3] - x1.size()[3]

        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
                        diffY // 2, diffY - diffY // 2])

        # for padding issues, see
        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd

        x = torch.cat([x2, x1], dim=1)

        x = self.conv(x)
        return x

class outconv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(outconv, self).__init__()
        self.conv = nn.Conv2d(in_ch, out_ch, 1)

    def forward(self, x):
        x = self.conv(x)
        return x


class GAU_1(nn.Module):
    def __init__(self, in_ch, out_ch, bilinear=True):
        super(GAU_1, self).__init__()
        self.up = nn.ConvTranspose2d(in_ch, in_ch, 2, stride=2)
        self.conv = double_conv(in_ch, out_ch)
        self.avgpool = nn.AdaptiveAvgPool2d(output_size=(1, 1))
        self.conv2d = nn.Conv2d(in_ch, out_ch, 3, padding=1)
        self.bn = nn.BatchNorm2d(out_ch)
        self.relu = nn.ReLU(inplace=True)
        self.bilinear = bilinear

    def forward(self, low, hight):
        if self.bilinear:
            hight = F.interpolate(hight, scale_factor=2, mode='bilinear', align_corners=True)
        else:
            hight = self.up(hight)

        # input is CHW
        diffY = low.size()[2] - hight.size()[2]
        diffX = low.size()[3] - hight.size()[3]

        hight = F.pad(hight, [diffX // 2, diffX - diffX // 2,
                        diffY // 2, diffY - diffY // 2])

        y = self.avgpool(hight)         #高级特征
        y = self.conv2d(y)
        y = self.bn(y)
        y = self.relu(y)
        y = torch.mul(low, y)
        output = torch.add(low, y)
        output = torch.add(hight, output)

        return output

class GAU_2(nn.Module):
    def __init__(self, in_ch, out_ch, bilinear=True):
        super(GAU_2, self).__init__()
        self.up = nn.ConvTranspose2d(in_ch, in_ch, 2, stride=2)
        self.conv = double_conv(in_ch, out_ch)
        self.avgpool = nn.AdaptiveAvgPool2d(output_size=(1, 1))
        self.conv2d = nn.Conv2d(in_ch, out_ch, 3, padding=1)
        self.bn = nn.BatchNorm2d(out_ch)
        self.relu = nn.ReLU(inplace=True)
        self.bilinear = bilinear

    def forward(self, low, hight):
        if self.bilinear:
            hight = F.interpolate(hight, scale_factor=2, mode='bilinear', align_corners=True)
        else:
            hight = self.up(hight)

        # input is CHW
        diffY = low.size()[2] - hight.size()[2]
        diffX = low.size()[3] - hight.size()[3]

        hight = F.pad(hight, [diffX // 2, diffX - diffX // 2,
                        diffY // 2, diffY - diffY // 2])

        y = self.avgpool(hight)         #高级特征
        y = self.conv2d(y)
        y = self.bn(y)
        y = self.relu(y)
        y = torch.mul(low, y)
        output = torch.add(low, y)

        hight = self.conv2d(hight)
        hight = self.bn(hight)
        hight = self.relu(hight)
        output = torch.add(hight, output)

        return output

class GAU_3(nn.Module):
    def __init__(self, in_ch, out_ch, bilinear=True):
        super(GAU_3, self).__init__()
        self.up = nn.ConvTranspose2d(in_ch, in_ch, 2, stride=2)
        self.conv = double_conv(in_ch, out_ch)
        self.avgpool = nn.AdaptiveAvgPool2d(output_size=(1, 1))
        self.conv2d = nn.Conv2d(in_ch, out_ch, 3, padding=1)
        self.bn = nn.BatchNorm2d(out_ch)
        self.relu = nn.ReLU(inplace=True)
        self.bilinear = bilinear

    def forward(self, low, hight):
        if self.bilinear:
            hight = F.interpolate(hight, scale_factor=2, mode='bilinear', align_corners=True)
        else:
            hight = self.up(hight)

        # input is CHW
        diffY = low.size()[2] - hight.size()[2]
        diffX = low.size()[3] - hight.size()[3]

        hight = F.pad(hight, [diffX // 2, diffX - diffX // 2,
                        diffY // 2, diffY - diffY // 2])

        y = self.avgpool(hight)         #高级特征
        y = self.conv2d(y)
        y = self.bn(y)
        y = self.relu(y)
        y = torch.mul(low, y)
        output = torch.add(low, y)
        hight = self.conv2d(hight)
        hight = self.bn(hight)
        hight = self.relu(hight)
        output = torch.add(hight, output)

        return output

class GAU_4(nn.Module):
    def __init__(self, in_ch, out_ch, bilinear=True):
        super(GAU_4, self).__init__()
        self.up = nn.ConvTranspose2d(in_ch, in_ch, 2, stride=2)
        self.conv = double_conv(in_ch, out_ch)
        self.avgpool = nn.AdaptiveAvgPool2d(output_size=(1, 1))
        self.conv2d = nn.Conv2d(in_ch, out_ch, 3, padding=1)
        self.bn = nn.BatchNorm2d(out_ch)
        self.relu = nn.ReLU(inplace=True)
        self.bilinear = bilinear

    def forward(self, low, hight):
        if self.bilinear:
            hight = F.interpolate(hight, scale_factor=2, mode='bilinear', align_corners=True)
        else:
            hight = self.up(hight)

        # input is CHW
        diffY = low.size()[2] - hight.size()[2]
        diffX = low.size()[3] - hight.size()[3]

        hight = F.pad(hight, [diffX // 2, diffX - diffX // 2,
                        diffY // 2, diffY - diffY // 2])

        y = self.avgpool(hight)         #高级特征
        y = self.conv2d(y)
        y = self.bn(y)
        y = self.relu(y)
        y = torch.mul(low, y)
        output = torch.add(low, y)
        hight = self.conv2d(hight)
        hight = self.bn(hight)
        hight = self.relu(hight)
        output = torch.add(hight, output)

        return output

class MANet(nn.Module):
    def __init__(self, classes):
        super(MANet, self).__init__()
        self.inc = inconv(3, 64)
        self.down1 = down(64, 128)
        self.down2 = down(128, 256)
        self.down3 = down(256, 512)
        self.down4 = down(512, 512)
        self.gau1 = GAU_1(512, 512)
        self.gau2 = GAU_2(512, 256)
        self.gau3 = GAU_3(256, 128)
        self.gau4 = GAU_4(128, 64)
        self.outc = outconv(64, classes)


    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.gau1(x4, x5)
        x = self.gau2(x3, x)
        x = self.gau3(x2, x)
        x = self.gau4(x1, x)
        out = self.outc(x)
        softmax = F.softmax(out, dim=1)
        out = torch.mul(out, softmax)

        return out




"""print layers and params of network"""
if __name__ == '__main__':
    # cuda
    # import os
    # os.environ["CUDA_VISIBLE_DEVICES"] = '0'
    # input = torch.Tensor(1,3,512,512).cuda()
    # model = MANet(classes=3).cuda()
    # model.eval()
    # print(model)
    # output = model(input)
    # print('MANet', output.size())
    # summary(model, (3, 512, 512))

    # cpu
    input = torch.Tensor(1, 3, 512, 512)
    model = MANet(classes=3)
    model.eval()
    print(model)
    output = model(input)
    print('MANet output size', output.size())
    summary(model, (3, 512, 512), device='cpu')
